import taichi as ti
import open3d as o3d
import numpy as np

from model_parser import mesh2pcd

ti.init(arch=ti.gpu)  # 使用GPU加速

@ti.kernel
def farthest_point_sampling(points: ti.template(), dists: ti.template(), sampled_indices: ti.template(), num_points: ti.template()):
    n = points.shape[0]

    seed = ti.random() * n
    sampled_indices[0] = int(seed)
    dists[0] = 0

    for i in range(1, num_points):
        best_index = -1
        best_dist = -1.0

        for j in range(n):
            if sampled_indices[j] < 0:
                continue

            p = points[sampled_indices[j]]
            
            for k in range(n):
                if sampled_indices[k] >= 0:
                    continue

                q = points[k]
                dist_sq = (p - q).norm_sqr()  # 计算点之间的欧氏距离的平方

                if dist_sq < dists[k]:
                    dists[k] = dist_sq

                if dist_sq > best_dist:
                    best_dist = dist_sq
                    best_index = k

        sampled_indices[i] = best_index


# 示例使用
# 在此处填充points数组，即将要进行最远点采样的点云数据

mesh = o3d.io.read_triangle_mesh("cache/Qma3crJ2cwrESHuXN6FiHySL3voCxy718y2sjg8rwv4qnc.obj")
if not mesh.has_vertex_normals():
    mesh.compute_vertex_normals()
vert = np.asarray(mesh.vertices)
norm = np.asarray(mesh.vertex_normals)

nsamples = 1000
points = ti.Vector.field(3 ,ti.f32, shape=(vert.shape[0]))  # 3维点的坐标
points.from_numpy(vert)
print(points)
dists = ti.field(ti.f32, shape=vert.shape[0])  # 存储距离的数组
sampled_indices = ti.field(ti.i32, shape=nsamples)  # 存储采样点索引的数组

farthest_point_sampling(points, dists, sampled_indices, nsamples)

vert = vert[sampled_indices[0],:]
norm = norm[sampled_indices[0],:]
print(sampled_indices)
print(vert)

# pcd = o3d.geometry.PointCloud()
# pcd.points = o3d.utility.Vector3dVector(points)
# pcd.normals = o3d.utility.Vector3dVector(norm)
